import numpy as np
import pandas as pd
import tushare as ts

from src.facility import *

# -------------------------------------------------

# for InsecureRequestWarning
urllib3.disable_warnings()
socket.setdefaulttimeout(30)

ts.set_token('fe53a8928a5442e82ec2ed835f02562c7f5ac24fe7edaa2a40299c8a')
pro = ts.pro_api()


# -------------------------------------------------


class StrategyCustom:
    # 手续费
    cost_sold = 0.0013
    cost_buy = 0.0003

    def __init__(self, begin_date, end_date, csv_path_pile_df, hold_day_count):
        self.begin_date = begin_date
        self.end_date = end_date
        self.capital_base = 100000
        self.hold_day_count = hold_day_count
        self.csv_path_pile_df = csv_path_pile_df
        print('init()-->', csv_path_pile_df)
        try:
            pile_df = pd.read_csv(csv_path_pile_df, encoding='utf-8')
            self.pile_df = pile_df[['买入日期', '卖出日期', '数量', '股票代码']]
            # print(self.pile_df)
        except FileNotFoundError:
            raise Exception('csv_path_pile_df?', csv_path_pile_df)
        pass

    # 今日收盘、今日开盘
    def buy_amount(yesterday_sum, amount, today_close, today_open):
        # print(yesterday_sum, amount, today_close, today_open)
        today_sum = yesterday_sum + today_close * amount - today_open * amount - today_open * amount * StrategyCustom.cost_buy
        result = (today_sum - yesterday_sum) / yesterday_sum
        return result, today_sum

    # 今日收盘、昨日收盘
    def hold_amount(yesterday_sum, amount, today_close, yesterday_close):
        today_sum = yesterday_sum + today_close * amount - yesterday_close * amount
        result = (today_sum - yesterday_sum) / yesterday_sum
        return result, today_sum

    # sold_price is 今日收盘（收盘卖） or 今日开盘（开盘卖），昨日收盘
    def sold_amount(yesterday_sum, amount, sold_price, yesterday_close):
        today_sum = yesterday_sum + sold_price * amount - sold_price * amount * StrategyCustom.cost_sold - yesterday_close * amount
        result = (today_sum - yesterday_sum) / yesterday_sum
        return result, today_sum

    # 只管修正
    def filter_pile_df(self):
        print('filter_pile_df()-->start...')
        for i in range(self.pile_df.index.size):
            buy_date = self.pile_df.at[i, '买入日期']
            sid = self.pile_df.at[i, '股票代码']
            # print('sid = [%s]' % (sid,))
            if pd.isna(sid):  # 这行没有股票代码，那么就什么都不计算
                continue
            sid = sid[:6] + ('.SZ' if sid[0] < '6' else '.SH')
            # print('sid = [%s]' % (sid,))

            df_stock = None  # df_stock是从0开始日期递减
            csv_path = 'csv/' + sid + '.csv'
            # print(csv_path)
            try:
                df_stock = pd.read_csv(csv_path, encoding='utf-8')
            except FileNotFoundError:
                df_stock = pro.daily(ts_code=sid)
                df_stock.to_csv(csv_path)
            df_stock = df_stock[['trade_date', 'close', 'open', 'pre_close']]  # 既是处理read后的空列，也是这样加快检索

            # 从小的日子开始，所以倒着df_stock，处理买入的情况
            j = df_stock.index.size - 1
            # print('外', i, j, buy_date, sid)
            for k in range(j, -1, -1):
                trade_date = df_stock.at[k, 'trade_date']
                trade_date = get_translated_trade_date(trade_date)
                # print('内', i, k, trade_date, buy_date)
                if trade_date == buy_date:  # 开市的买入
                    sold_index = k - self.hold_day_count
                    if sold_index < 0:
                        self.pile_df.at[i, '卖出日期'] = ''
                    else:
                        sold_date = df_stock.at[sold_index, 'trade_date']  # 下标从大到小，对应日期递增，所以减
                        sold_date = get_translated_trade_date(sold_date)
                        self.pile_df.at[i, '卖出日期'] = sold_date
                        # print('内', i, k, trade_date, buy_date, sold_date)
                    break
                if trade_date > buy_date:  # 修正未开市的买入
                    self.pile_df.at[i, '买入日期'] = trade_date
                    sold_index = k - self.hold_day_count
                    if sold_index < 0:
                        self.pile_df.at[i, '卖出日期'] = ''
                    else:
                        sold_date = df_stock.at[sold_index, 'trade_date']
                        sold_date = get_translated_trade_date(sold_date)
                        # print(sold_date)
                        self.pile_df.at[i, '卖出日期'] = sold_date
                        # print('内', i, k, trade_date, buy_date, sold_date)
                    break
        self.pile_df.to_csv(self.csv_path_pile_df)
        print('filter_pile_df()-->over...')
        pass

    # 这里采用总资产计算公式1
    # 指定begin_date和 end_date
    def get_df_result_1(self):
        original_capital = self.capital_base
        cum_capital = original_capital

        # print('self.get_df_result_1()-->', self.pile_df)
        # self.pile_df按从0开始，日期递增
        for i in range(self.pile_df.index.size):
            buy_date = self.pile_df.at[i, '买入日期']
            sold_date = self.pile_df.at[i, '卖出日期']
            sid = self.pile_df.at[i, '股票代码']
            # print('sid = [%s]' % (sid,))
            if pd.isna(sid):  # 这行没有股票代码，那么就什么都不计算
                continue
            sid = sid[:6] + ('.SZ' if sid[0] < '6' else '.SH')
            # print('sid = [%s]' % (sid,))
            amount = self.pile_df.at[i, '数量']

            # 买入日期超过结束日期退出，自然卖出日期更超过，而且之后的记录是越来越大
            if pd.notna(buy_date) and buy_date > self.end_date:  # buy_date有NAN，是由于set_pile_df()中的无头卖出
                # print('超过今天才退出', buy_date, end_date)
                break
            try:
                # 卖出日期小于起始日期不处理，自然买入日期更小于
                if pd.notna(sold_date) and sold_date < self.begin_date:  # sold_date 有NAN，是由于set_pile_df()中的无尾买入
                    # print('小于买入日期不处理', sold_date, begin_date)
                    continue
            except TypeError:
                raise Exception('?', sold_date, self.begin_date)

            select_flag_lists = ['buy', 'hold', 'sold']
            if pd.notna(buy_date) and buy_date == self.begin_date:  # 起始日期的股票持仓、卖出不算
                continue
            if pd.notna(sold_date) and sold_date == self.end_date:
                select_flag_lists = ['hold', 'sold']

            df_stock = None  # df_stock是从0开始日期递减
            csv_path = 'csv/' + sid + '.csv'
            # print(csv_path)
            try:
                df_stock = pd.read_csv(csv_path, encoding='utf-8')
            except FileNotFoundError:
                df_stock = pro.daily(ts_code=sid)
                df_stock.to_csv(csv_path)
            df_stock = df_stock[['trade_date', 'close', 'open', 'pre_close']]  # 既是处理read后的空列，也是这样加快检索

            # 从小的日子开始，所以倒着df_stock，处理买入的情况
            j = df_stock.index.size - 1
            for k in range(j, -1, -1):
                trade_date = df_stock.at[k, 'trade_date']
                trade_date = get_translated_trade_date(trade_date)
                # print('?', k, trade_date, self.pile_df.at[i, '买入日期'])

                # 处理停市的买的日子
                if trade_date == buy_date:  # 有交易情况
                    j = k
                    if buy_date < self.begin_date:  # 买入日期小于起始不处理买入，然后去看是否处理持仓和卖出
                        break
                    if 'buy' in select_flag_lists:
                        today_close = df_stock.at[k, 'close']
                        today_open = df_stock.at[k, 'open']
                        cum_capital = StrategyCustom.buy_amount(yesterday_sum=cum_capital,
                                                                amount=amount,
                                                                today_close=today_close,
                                                                today_open=today_open)[1]
                        print('买', k, trade_date, amount, '收', today_close, '开', today_open, cum_capital)
                    break

            # 处理持仓和卖出
            if pd.notna(sold_date):
                for k in range(j - 1, -1, -1):  # j-1自然因为买入卖出不可能在同一天
                    trade_date = df_stock.at[k, 'trade_date']
                    trade_date = get_translated_trade_date(trade_date)

                    # 交易日期小于起始日期，就去下一个交易日
                    if trade_date < self.begin_date:
                        continue
                    # 交易日期超过结束日期，就不处理持仓和卖了
                    if trade_date > self.end_date:
                        break

                    # 处理停市的卖的日子：
                    if trade_date == sold_date:
                        if 'sold' in select_flag_lists:
                            sold_price = df_stock.at[k, 'open']
                            yesterday_close = df_stock.at[k, 'pre_close']
                            cum_capital = StrategyCustom.sold_amount(yesterday_sum=cum_capital,
                                                                     amount=amount,
                                                                     sold_price=sold_price,
                                                                     yesterday_close=yesterday_close)[1]
                            print('卖', k, trade_date, amount, '今开', sold_price, '昨收', yesterday_close, cum_capital)
                        break

                    # 处理持仓
                    if 'hold' in select_flag_lists:
                        today_close = df_stock.at[k, 'close']
                        yesterday_close = df_stock.at[k, 'pre_close']
                        cum_capital = StrategyCustom.hold_amount(yesterday_sum=cum_capital,
                                                                 amount=amount,
                                                                 today_close=today_close,
                                                                 yesterday_close=yesterday_close)[1]
                        print('持', k, trade_date, amount, '今收', today_close, '昨收', yesterday_close, cum_capital)

            # print('\r [%d]/[%d]...' % (i, self.pile_df.index.size), end=' ')
        # print('')

        result = (cum_capital - original_capital) / original_capital
        return result


def back_test_one(begin_date, end_date, csv_path_pile_df, hold_day_count):
    my_strategy = StrategyCustom(
        begin_date=begin_date,
        end_date=end_date,
        csv_path_pile_df=csv_path_pile_df,
        hold_day_count=hold_day_count
    )
    my_strategy.filter_pile_df()
    result = my_strategy.get_df_result_1()
    print(result)
    return result


def litter_step(hold_day_count, days_count, end_date):
    csv_path_strategy_lists = 'csv/pool' + '.csv'
    df_strategy_lists = pd.read_csv(csv_path_strategy_lists, encoding='utf-8')
    df_strategy_lists = df_strategy_lists[
        ['unique_id', 'id', '策略作者', '策略名字',
         '胜率', '最大回撤', '夏普率', '运行交易日',
         '最近1周收益率', '300基准收益率',
         '持有一天的累计一周收益', '持有两天的累计一周收益', '持有一天的累计一月收益', '持有两天的累计一月收益']
    ]

    begin_date = get_day(end_date, -days_count)
    operation_num = hold_day_count + days_count
    operation_dict = {
        8: '持有一天的累计一周收益',
        9: '持有两天的累计一周收益',
        31: '持有一天的累计一月收益',
        32: '持有两天的累计一月收益'
    }

    result_lists = []
    for i in range(df_strategy_lists.index.size):
        csv_path_pile_df = 'csv/D1_' + df_strategy_lists.at[i, 'unique_id'] + '.csv'
        result = back_test_one(begin_date, end_date, csv_path_pile_df, hold_day_count)
        df_strategy_lists.at[i, operation_dict.get(operation_num)] = result
        # break

    df_strategy_lists.to_csv(csv_path_strategy_lists)
    pass


def main():
    end_date = '2021-01-29'
    litter_step(1, 7, end_date=end_date)
    litter_step(2, 7, end_date=end_date)
    litter_step(1, 30, end_date=end_date)
    litter_step(2, 30, end_date=end_date)
    pass


if __name__ == '__main__':
    main()
